﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;

namespace AntlrTS.TypeInference
{
    class InfereceHelper
    {
        private static TSTerm tagImpl(TSTerm form, TSType ty)
        {
            if (form is Annotate)
            {
                return new Annotate((form as Annotate).body, ty);
            }
            return new Annotate(form, ty);
        }

        private static TSTerm constructorImpl(TSTerm form,string x,Func<TSTerm, TSTerm> coercion,bool isCtor=true)
        {
            var temp = new Var("?");
            var coerciedTmp = coercion(temp);
            string para1 = "?CTOR";
            string para2 = "?ARG";
            string paratmp = "";
            if (coerciedTmp is Annotate)
            {
                if (!isCtor)
                {
                    paratmp = para1;
                    para1 = para2;
                    para2 = paratmp;
                }
                if (form is Annotate)
                {
                    var tmp=new Dictionary<string,TSType>();
                    tmp.Add(para1+x,(coerciedTmp as Annotate).type);
                    return new Annotate((form as Annotate).body,
                        (form as Annotate).type.subst(tmp));
                }
                return new Annotate(form,
                    new Composite((coerciedTmp as Annotate).type,
                        new Slot(para2+ x)
                        ,false)
                    );
            }
            return coerciedTmp.subst("?", form);
        }

        public static string rawNameOfskolmeizedName(string n)
        {
            Regex re = new Regex(@"\.\d+\.");
            return re.Replace(n, "");
        }

        public static string rawNameToskolmeizedName(string u,string n) 
        {
	        return "." + u + "." + n;
        }

        public static Func<TSTerm, TSTerm> tag(TSType ty)
        {
            return term => tagImpl(term, ty);
        }

        public static Func<TSTerm, TSTerm> constructor(string x, Func<TSTerm, TSTerm> coercion)
        {
            return y => constructorImpl(y, x, coercion);
        }

        public static Func<TSTerm, TSTerm> arg(string x, Func<TSTerm, TSTerm> coercion)
        {
            return y => constructorImpl(y, x, coercion,false);
        }

        //public static buildGL(Dictionary<string,Slot> m1,TSTerm e)
        //{
        //    if (e is Generic)
        //    {
        //        var m2=new HashSet<TSType>((e as Generic).quantifiers);
        //        foreach (var item in m1.Values)
        //        {
        //             m2.Add(item);
        //        }
        //        return new Generic(m2.ToList(),(e as Generic).body);
        //    }
        //}

        public static string generateBinder(long nRef, HashSet<string> used)
        {
            nRef += 1;
            string name = "t" + nRef.ToString();
            while (used.Contains(name))
            {
                nRef += 1;
                name = "t" + nRef.ToString();
            }
            return name;
        }
    }


    class TSType
    {
        public TSType instancetype = null;

        //public bool 

        public TSType() { }

        public virtual List<string> format() { return new List<string>(); }

        public virtual Dictionary<int, MetaSlot> getMetaSlots()
        {
            var a = new Dictionary<int, MetaSlot>();
            getMetaSlotsImpl(a);
            return new Dictionary<int, MetaSlot>();
        }

        public virtual void getMetaSlotsImpl(Dictionary<int, MetaSlot> a)
        {
            return;
        }

        public virtual HashSet<string> getFreeSlots()
        {
            var a = new HashSet<string>();
            var bound = new HashSet<string>();
            this.getFreeSlotsImpl(bound, a);
            return new HashSet<string>();
        }

        public virtual void getFreeSlotsImpl(HashSet<string> bound, HashSet<string> a)
        {
            return;
        }

        public virtual HashSet<string> getBinders()
        {
            var a = new HashSet<string>();
            this.getBindersImpl(a);
            return a;
        }

        public virtual void getBindersImpl(HashSet<string> a)
        {
            return;
        }

        public virtual TSType subst(Dictionary<string, TSType> variables)
        {
            return this;
        }

        public virtual Tuple<TSType, Func<TSTerm, TSTerm>> instantiate(TSEnviroment te)
        {
            return new Tuple<TSType, Func<TSTerm, TSTerm>>(this, InfereceHelper.tag(this));
        }

        public virtual Tuple<Dictionary<string, Slot>, TSType, Func<TSTerm, TSTerm>> skolmeize(TSEnviroment te)
        {
            return new Tuple<Dictionary<string, Slot>, TSType, Func<TSTerm, TSTerm>>(
                new Dictionary<string, Slot>(),
                this,
                InfereceHelper.tag(this)
                );
        }

        public virtual Forall generalize(TSEnviroment te, IEnumerable<MetaSlotVal> mvs)
        {
            var binders = getBinders();
            long nRef = 0;
            var newBinders = new List<Slot>();
            var forallStrList = new List<string>();
            foreach (var item in mvs)
            {
                var newBinder = new Slot(InfereceHelper.generateBinder(nRef, binders));
                item.typeRef = newBinder;
                newBinders.Add(newBinder);
                forallStrList.Add(newBinder.name);
            }
            return new Forall(forallStrList, this.zonk());
        }

        public virtual TSType zonk()
        {
            return this;
        }

        public virtual Tuple<TSType, Func<TSTerm, TSTerm>> instSigmaInfer(TSEnviroment te)
        {
            return instantiate(te);
        }
    }

    class Primitive : TSType
    {
        public string name {get; private set; }
        public Primitive(string name)
            : base()
        {
            this.name = name;
        }

        public override List<string> format()
        {
            var x = base.format();
            x.Add(name);
            return x;
        }
    }



    class MetaSlot : TSType
    {
        public  MetaSlotVal arg;
        public MetaSlot(MetaSlotVal arg)
        {
            this.arg = arg;
        }

        public override void getMetaSlotsImpl(Dictionary<int, MetaSlot> a)
        {
            /*refactor*/
            if (!a.ContainsKey((int)arg.id))
            {
                a.Add((int)arg.id, new MetaSlot(arg));
            }
        }

        public override TSType zonk()
        {
            if (this.arg != null && this.arg.typeRef != null)
            {
                var t1 = this.arg.typeRef.zonk();
                this.arg.typeRef = t1;
                return t1;
            }
            return this;
        }
    }

    class Slot : TSType
    {
        public string name;
        public Slot(string name)
            : base()
        {
            this.name = name;
        }

        public override void getFreeSlotsImpl(HashSet<string> bound, HashSet<string> a)
        {
            if (!bound.Contains(name) && !a.Contains(name))
            {
                a.Add(name);
            }
        }

        public bool isSkolem()
        {
            return this.name[0] == '.';
        }

        public override TSType subst(Dictionary<string, TSType> variables)
        {
            if (variables.ContainsKey(name))
            {
                return variables[name];
            }
            return this;
        }

        public bool equalTo(Slot s)
        {
            if (this.isSkolem() && s.isSkolem())
            {
                return InfereceHelper.rawNameOfskolmeizedName(this.name) == InfereceHelper.rawNameOfskolmeizedName(s.name);
            }
            else if (!this.isSkolem() && !s.isSkolem())
            {
                return this.name == s.name;
            }
            return false;
        }
    }

    class MetaSlotVal
    {

        public long id;
        public TSType typeRef;
        public MetaSlotVal(long id, TSType typeRef)
        {
            this.id = id;
            this.typeRef = typeRef;
        }

        public bool equalTo(MetaSlotVal ms)
        {
            return ms != null && ms is MetaSlotVal && ms.id == this.id;
        }
    }

    class Composite : TSType
    {
        bool contravariant;

        public TSType fn;

        public TSType arg;
        public Composite(TSType fn, TSType arg, bool contravariant)
            : base()
        {
            this.fn = fn;
            this.arg = arg;
            this.contravariant = contravariant;
        }

        public override void getFreeSlotsImpl(HashSet<string> bound, HashSet<string> a)
        {
            fn.getFreeSlotsImpl(bound, a);
            arg.getFreeSlotsImpl(bound, a);
        }

        public override void getBindersImpl(HashSet<string> a)
        {
            fn.getBindersImpl(a);
            arg.getBindersImpl(a);
        }

        public override void getMetaSlotsImpl(Dictionary<int, MetaSlot> a)
        {
            fn.getMetaSlotsImpl(a);
            arg.getMetaSlotsImpl(a);
        }

        public override TSType subst(Dictionary<string, TSType> m)
        {
            return new Composite(this.fn.subst(m), this.arg.subst(m), this.contravariant);
        }

        //public override Tuple<Dictionary<string, Slot>, TSType, Func<TSTerm, TSTerm>> skolmeize(TSEnviroment te)
        //{
        //    Tuple<Dictionary<string, Slot>, TSType, Func<TSTerm, TSTerm>> tempTuple = this.fn.skolmeize(te);
        //    if (this.contravariant)
        //    {
        //        var argtype = this.arg;
        //        var skolRho=
        //    }
        //    return base.skolmeize(te);
        //}

        public override TSType zonk()
        {
            return new Composite(this.fn.zonk(), this.arg.zonk(), this.contravariant);
        }
    }

    class Forall : TSType
    {
        List<string> quantifiers;
        TSType body;
        public Forall(List<string> quantifiers, TSType body)
        {
            this.quantifiers = quantifiers;
            this.body = body;
        }

        public override void getMetaSlotsImpl(Dictionary<int, MetaSlot> a)
        {
            body.getMetaSlotsImpl(a);
        }

        public override void getFreeSlotsImpl(HashSet<string> bound, HashSet<string> a)
        {
            var b1 = new HashSet<string>(bound);
            foreach (var item in quantifiers)
            {
                b1.Add(item);
            }
            body.getFreeSlotsImpl(b1, a);
        }

        public override void getBindersImpl(HashSet<string> a)
        {
            foreach (var item in quantifiers)
            {
                a.Add(item);
            }
            body.getBindersImpl(a);
        }

        public override TSType subst(Dictionary<string, TSType> variables)
        {
            var m1 = new Dictionary<string, TSType>(variables);
            foreach (var item in quantifiers)
            {
                m1.Remove(item);
            }
            return new Forall(quantifiers, body.subst(m1));
        }

        //public override Tuple<TSType, Func<Term, Term>> instantiate(TSEnviroment te)
        //{
        //    var m = new Dictionary<string, MetaSlot>();
        //    foreach (var item in quantifiers)
        //    {
        //        m.Add(item, new MetaSlot(te.newMetaSlotVal()));
        //    }
        //    /*todo*/
        //    return new Tuple<TSType, Func<Term, Term>>(
        //        body.subst(m),
        //        x=>
        //        {
        //            if (m.Count>0)
        //            {
        //                if (x is Invoke && x.f)
        //                {

        //                }
        //            }
        //        }
        //        );
        //}

        public override Tuple<Dictionary<string, Slot>, TSType, Func<TSTerm, TSTerm>> skolmeize(TSEnviroment te)
        {
            /*todo*/
            //var m=new Dictionary
            return base.skolmeize(te);
        }
    }

    class TSEnviroment
    {
        private long uniqs;
        private Dictionary<string, TSType> variables;
        public TSEnviroment(long uniqs, Dictionary<string, TSType> variables)
        {
            this.uniqs = uniqs;
            this.variables = variables;
        }

        public TSEnviroment extend(string name, TSType type)
        {
            var newVariables = new Dictionary<string, TSType>(this.variables);
            newVariables.Add(name, type);
            return new TSEnviroment(this.uniqs, newVariables);
        }

        public TSEnviroment extendN(IEnumerable<KeyValuePair<string, TSType>> terms)
        {
            var newVariables = new Dictionary<string, TSType>(this.variables);
            foreach (var item in terms)
            {
                newVariables.Add(item.Key, item.Value);
            }
            return new TSEnviroment(this.uniqs, newVariables);
        }

        public TSType find(string name)
        {
            if (variables.ContainsKey(name))
            {
                return variables[name];
            }
            throw new Exception("can't find type");
        }

        public long newUniqs()
        {
            this.uniqs += 1;
            return uniqs;
        }

        public string newVar(string tag = "")
        {
            this.uniqs += 1;
            return "TS_" + tag + "_" + this.uniqs.ToString();
        }

        public MetaSlotVal newMetaSlotVal()
        {
            return new MetaSlotVal(newUniqs(), null);
        }

        public string newSkolemVariable(string n)
        {
            return InfereceHelper.rawNameToskolmeizedName(this.newUniqs().ToString(), n);
        }

        public List<TSType> getTypes()
        {
            return variables.Values.ToList();
        }

        public IEnumerable getMetaSlotVars(IEnumerable<TSType> tys)
        {
            foreach (var item in tys)
            {
                var ty = item.zonk();
                yield return ty.getFreeSlots();
            }
        }
    }

    class TSTerm
    {
        public TSTerm() { }

        public virtual bool isAtomic(){return false;}

        public virtual TSTerm subst(string id, TSTerm replacement) { return this; }
    }

    class Literal:TSTerm
    {
        public Object lit { get; private set; }
        public Literal(Object lit):base()
        {
            this.lit = lit;
        }

        public override bool isAtomic()
        {
            return true;
        }
    }

    class Var:TSTerm
    {
        public string name { get; private set; }

        public Var(string name):base()
        {
            this.name = name;
        }

        public override bool isAtomic()
        {
            return true;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            if (id==name)
            {
                return replacement; 
            }
            return this;
        }
    }

    class FuncApply:TSTerm
    {
        public TSTerm fn { get; private set; }
        public TSTerm arg { get; private set; }
        public FuncApply(TSTerm fn,TSTerm arg):base()
        {
            this.fn = fn;
            this.arg = arg;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            return new FuncApply(this.fn.subst(id, replacement), this.arg.subst(id, replacement));
        }
    }

    class Lambda:TSTerm
    {
        public string param { get; private set; }
        public TSTerm body { get; private set; }
        public Lambda(string param,TSTerm body):base()
        {
            this.param = param;
            this.body = body;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            if (id!=this.param)
            {
                return new Lambda(this.param, this.body.subst(id, replacement));
            }
            return this;
        }
    }

    class AnnoatedLambda:TSTerm
    {
        private string param;
        private TSType type;
        private TSTerm body;
        public AnnoatedLambda(string param,TSType type,TSTerm body):base()
        {
            this.param = param;
            this.type = type;
            this.body = body;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
           if (id!=this.param)
            {
                return new AnnoatedLambda(this.param,type, this.body.subst(id, replacement));
            }
           return this;
        }
    }

    class LetBind:TSTerm
    {
        public TSTerm body { get; private set; }
        public Dictionary<string, TSTerm> terms { get; private set; }
        public LetBind(Dictionary<string,TSTerm>terms,TSTerm body):base()
        {
            this.terms = terms;
            this.body = body;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            var terms1 = new Dictionary<string, TSTerm>();
            bool overrided = false;
            foreach (var item in terms)
            {
                if (item.Key==id)
                {
                    overrided = true;
                }
                terms1.Add(item.Key, item.Value.subst(id, replacement));
            }
            return new LetBind(terms1, !overrided ? this.body.subst(id, replacement) : this.body);
        }
    }

    class LetRecBind:TSTerm
    {
        private List<Tuple<string, TSTerm, TSType>> terms;
        private TSTerm body;
        public LetRecBind(List<Tuple<string, TSTerm, TSType>> terms, TSTerm body)
            : base()
        {
            this.terms = terms;
            this.body = body;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            var terms1 = new List<Tuple<string, TSTerm, TSType>>();
            bool overrided = false;
            foreach (var item in terms)
            {
                if (item.Item1  == id)
                {
                    overrided = true;
                }
            }
            foreach (var item in terms)
            {
                if (overrided)
                {
                    terms1.Add(item);
                }
                else
                {
                    terms1.Add(new Tuple<string, TSTerm, TSType>(
                        item.Item1,
                        item.Item2.subst(id, replacement),
                        item.Item3
                        ));
                }
            }
            return new LetRecBind(terms1, !overrided ? this.body.subst(id, replacement) : this.body);
        }
    }

    class Annotate:TSTerm 
    {
        public TSType type;
        public TSTerm body;
        public Annotate(TSTerm body,TSType type):base()
        {
            this.body = body;
            this.type = type;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            return new Annotate(this.body.subst(id, replacement), this.type);
        }
    }

    class Generic:TSTerm
    {
        public List<Slot> quantifiers;
        public TSTerm body;
        public Generic(List<Slot> quantifiers,TSTerm body):base()
        {
            this.quantifiers = quantifiers;
            this.body = body;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            return new Generic(this.quantifiers, this.body.subst(id, replacement));
        }
    }

    class Inst:TSTerm
    {
        private Dictionary<string, TSType> args;
        private TSTerm body;
        public Inst(TSTerm body,Dictionary<string,TSType> args):base()
        {
            this.body = body;
            this.args = args;
        }

        public override TSTerm subst(string id, TSTerm replacement)
        {
            return new Inst(this.body.subst(id, replacement), this.args);
        }
    }

    public class TestHM
    {

        private static bool badType(TSType t)
        {
            return t is Slot && !(t as Slot).isSkolem();
        }

        private static void Unify(TSType t1,TSType t2)
        {
            if (badType(t1)||badType(t2))
            {
                throw new Exception("type t1,t2 are slot but both them are not skolem");
            }
            if (t1 is Slot && t2 is Slot && (t1 as Slot).equalTo(t2 as Slot))
            {
                return;
            }
            if (t1 is MetaSlot && t2 is MetaSlot && (t1 as MetaSlot).arg.equalTo((t2 as MetaSlot).arg))
            {
                return;
            }
            if (t1 is MetaSlot)
            {
                UnifyMetaSlot((t1 as MetaSlot).arg, t2);
                return;
            }
            if (t2 is MetaSlot)
            {
                UnifyMetaSlot((t2 as MetaSlot).arg, t1);
                return;
            }
            if (t1 is Composite && t2 is Composite)
            {
                Unify((t1 as Composite).fn, (t2 as Composite).fn);
                Unify((t1 as Composite).arg, (t2 as Composite).arg);
                return;
            }
            if (t1 is Primitive && t2 is Primitive && (t1 as Primitive).name==(t2 as Primitive).name)
            {
                return;
            }
            throw new Exception("can not unify");
        }

        private static void UnifyMetaSlot(MetaSlotVal msv,TSType t2)
        {
            if (msv.typeRef!=null)
            {
                Unify(msv.typeRef, t2);
            }
            else
            {
                UnifyUnbound(msv, t2);
            }
        }

        private static void UnifyUnbound(MetaSlotVal msv,TSType t2)
        {
            if (t2 is MetaSlot)
            {
                var msv2 = (t2 as MetaSlot).arg;
                if (msv2.typeRef!=null)
                {
                    Unify(new MetaSlot(msv), msv2.typeRef);
                    return;
                }
                else
                {
                    msv.typeRef = t2;
                    return;
                }
            }
            else
            {
                var msvs2 = t2.getMetaSlots();
                if (msvs2.ContainsKey((int)msv.id))
                {
                    /*refactor*/
                    throw new Exception("recursive type");
                }
                else
                {
                    msv.typeRef = t2;
                    return;
                }
            }
        }

        private static TSType Analyze(TSTerm t,TSEnviroment te)
        {
            if(t is Var)
            {
                return te.find((t as Var).name);
            }
            else if(t is Literal)
            {
                return new Primitive((t as Literal).lit.GetType().ToString());
            }
            else if(t is FuncApply)
            {
                var funcType = Analyze((t as FuncApply).fn,te);
                var argType = Analyze((t as FuncApply).arg,te);
                var resulttype = new MetaSlot(te.newMetaSlotVal());
                Unify(new Composite(argType, resulttype, false), funcType);
                return resulttype.arg.typeRef;
            }
            else if(t is Lambda)
            {
                var arg_type = new MetaSlot(te.newMetaSlotVal());
                var teCopy = te.extend((t as Lambda).param, arg_type);
                var result_type = Analyze((t as Lambda).body,teCopy);
                return new Composite(arg_type.arg.typeRef, result_type, false);
            }
            else if(t is LetBind)
            {
                Dictionary<string, TSType> tmp = new Dictionary<string, TSType>();
                foreach (var item in (t as LetBind).terms)
                {
                    tmp.Add(item.Key, Analyze(item.Value, te));
                }
                var teCopy = te.extendN(tmp);
                var defn_type = Analyze((t as LetBind).body, teCopy);
                return defn_type;
            }
            return new TSType();
        }

        public static void testAnalyze()
        {
            TSEnviroment te = new TSEnviroment(0, new Dictionary<string, TSType>());
            //te=te.extend("f", new Primitive("int"));
            te = te.extend("f", new Composite(new Composite(new Primitive("bool"), new Primitive("bool"), false), new Primitive("bool"), false));
            
            te=te.extend("k", new Primitive("bool"));
            //var lettest=new LetBind()
            var test = new FuncApply(new Var("f"), new Var("k"));
            var res_type=Analyze(test, te);
        }
    }
}
